"""根据 VQModel 的 base.yaml 自动生成带 clip 版本的 config
1. 加载指定的 yaml 文件
2. 添加 clip 相关的参数 (可能用到原 yaml 中的值)
3. 保存新 config 到指定路径 (一般和原 yaml 文件在同一文件夹下，命名是 clip specific 的)
"""
# TODO : 添加指定 ckpt_path 功能
import yaml
from omegaconf import OmegaConf
from omegaconf.dictconfig import DictConfig
from pathlib import Path

from argparse import ArgumentParser

def prepare_args():
    parser=ArgumentParser()
    parser.add_argument(
        "--base_path",
        type=str,
    )
    parser.add_argument(
        "--clip_res",
        type=int,
        help="feature resolution to do multi head attn pooling",
        default=8,
    )
    parser.add_argument(
        "--clip_feat_dim",
        type=int,
        help="depending on which CLIP model you use (default ViT-B/16)",
        default=512,
    )
    parser.add_argument(
        "--num_attn_head",
        type=int,
        default=8,
    )
    parser.add_argument(
        "--clip_pos",
        type=str,
        help="['pre', 'post']. add clip loss before or after VQ action. "
    )
    parser.add_argument(
        "--clip_weight",
        type=float,
        default=1.0,
    )
    parser.add_argument(
        "--clip_model_name",
        type=str,
        help="OpenAI CLIP model name",
        default="ViT-B-16",
    )
    parser.add_argument(
        "--train_clip_list_file",
        type=str,
        default="data/OpenImages/custom_train_clip_1743042.txt",
    )
    parser.add_argument(
        "--val_clip_list_file",
        type=str,
        default="data/OpenImages/custom_val_clip_41620.txt",
    )
    args=parser.parse_args()
    return args

if __name__ == "__main__":
    args=prepare_args()
    
    # # ! debug
    # ##########################
    # args.base_path="configs/vq-f4/base.yaml"
    # args.clip_pos='pre'
    # ##########################

    base_cfg=OmegaConf.load(args.base_path)
    res = base_cfg.model.params.ddconfig.resolution
    down_f = 2**(len(base_cfg.model.params.ddconfig.ch_mult)-1)
    
    base_cfg.model.target = "vae.models.autoencoder.VQModelWithCLIP"
    
    assert res % down_f ==0
    assert args.clip_pos in ['pre','post']
    
    clip_cfg = DictConfig({
        "in_res": res // down_f,
        "f_res": args.clip_res,
        "in_dim": base_cfg.model.params.ddconfig.z_channels,
        "clip_feature_dim": args.clip_feat_dim,
        "heads": args.num_attn_head,
        "clip_pre_vq": True if args.clip_pos=='pre' else False,
    })
    
    base_cfg.model.params.clip_cfg=clip_cfg
    
    # TODO: clip_weight, data 中 with_clip_feat
    base_cfg.model.params.lossconfig.clip_weight = args.clip_weight
    base_cfg.data.params.train.params.clip_list_file = args.train_clip_list_file
    base_cfg.data.params.validation.params.clip_list_file = args.val_clip_list_file
    
    
    save_path = Path(args.base_path).parent / f'CLIP{args.clip_pos}VQ_{args.clip_model_name}.yaml'
    OmegaConf.save(base_cfg, str(save_path))